#include <gtest/gtest.h>
#include "viml/lexer.h"
#include "viml/parser.h"
#include "viml/ast.h"

namespace viml {

class CompilerBasicTest : public ::testing::Test {
protected:
    void SetUp() override {}
    void TearDown() override {}
};

class ParserBasicTest : public ::testing::Test {
protected:
    void SetUp() override {}
    void TearDown() override {}
};

// Additional Lexer Tests for new features
TEST_F(CompilerBasicTest, CommentInString) {
    viml::Lexer lexer("let msg = \"Hello \\\"world\\\" test\"");
    auto tokens = lexer.tokenize();
    
    // 应该有5个token: LET, IDENTIFIER(msg), ASSIGN, STRING, EOF
    EXPECT_EQ(tokens.size(), 5);
    EXPECT_EQ(tokens[3].type, viml::TokenType::STRING);
    EXPECT_EQ(tokens[3].value, "Hello \"world\" test");
}

TEST_F(CompilerBasicTest, NewlineSeparatedStatements) {
    viml::Lexer lexer("let x = 1\nlet y = 2");
    auto tokens = lexer.tokenize();
    
    // Should have NEWLINE token between statements
    bool found_newline = false;
    for (const auto& token : tokens) {
        if (token.type == viml::TokenType::NEWLINE) {
            found_newline = true;
            break;
        }
    }
    EXPECT_TRUE(found_newline);
}

TEST_F(CompilerBasicTest, DictLiteralTokens) {
    viml::Lexer lexer("{\"key\": \"value\"}");
    auto tokens = lexer.tokenize();
    
    EXPECT_EQ(tokens[0].type, viml::TokenType::LBRACE);
    EXPECT_EQ(tokens[1].type, viml::TokenType::STRING);  // "key"
    EXPECT_EQ(tokens[2].type, viml::TokenType::COLON);
    EXPECT_EQ(tokens[3].type, viml::TokenType::STRING);  // "value"
    EXPECT_EQ(tokens[4].type, viml::TokenType::RBRACE);
    EXPECT_EQ(tokens[5].type, viml::TokenType::END_OF_FILE);
}

// Parser Tests for new features
TEST_F(ParserBasicTest, SimpleLetStatement) {
    viml::Parser parser("let x = 42");
    auto program = parser.parse();
    
    EXPECT_EQ(program->statements.size(), 1);
    
    auto let_stmt = dynamic_cast<viml::LetStatement*>(program->statements[0].get());
    EXPECT_NE(let_stmt, nullptr);
    EXPECT_EQ(let_stmt->name, "x");
    
    auto number_literal = dynamic_cast<viml::NumberLiteral*>(let_stmt->value.get());
    EXPECT_NE(number_literal, nullptr);
    EXPECT_EQ(number_literal->value, 42.0);
}

TEST_F(ParserBasicTest, BinaryExpression) {
    viml::Parser parser("let result = a + b * c");
    auto program = parser.parse();
    
    auto let_stmt = dynamic_cast<viml::LetStatement*>(program->statements[0].get());
    EXPECT_NE(let_stmt, nullptr);
    
    auto binary_expr = dynamic_cast<viml::BinaryExpression*>(let_stmt->value.get());
    EXPECT_NE(binary_expr, nullptr);
    EXPECT_EQ(binary_expr->op, viml::BinaryExpression::Operator::ADD);
}

TEST_F(ParserBasicTest, IfStatement) {
    viml::Parser parser("if x > 5\necho \"greater\"\nendif");
    auto program = parser.parse();
    
    EXPECT_EQ(program->statements.size(), 1);
    
    auto if_stmt = dynamic_cast<viml::IfStatement*>(program->statements[0].get());
    EXPECT_NE(if_stmt, nullptr);
    EXPECT_EQ(if_stmt->then_body.size(), 1);  // 现在应该是1
    
    auto binary_expr = dynamic_cast<viml::BinaryExpression*>(if_stmt->condition.get());
    EXPECT_NE(binary_expr, nullptr);
    EXPECT_EQ(binary_expr->op, viml::BinaryExpression::Operator::GT);
}

TEST_F(ParserBasicTest, IfElseIfStatement) {
    viml::Parser parser("if x < 5\necho \"less\"\nelseif x == 5\necho \"equal\"\nelse\necho \"greater\"\nendif");
    auto program = parser.parse();
    
    auto if_stmt = dynamic_cast<viml::IfStatement*>(program->statements[0].get());
    EXPECT_NE(if_stmt, nullptr);
    EXPECT_EQ(if_stmt->then_body.size(), 1);
    // elseif被转换为嵌套的if语句，所以else_body应该有1个元素
    EXPECT_EQ(if_stmt->else_body.size(), 1);
    
    // elseif是转换为嵌套的if
    auto nested_if = dynamic_cast<viml::IfStatement*>(if_stmt->else_body[0].get());
    EXPECT_NE(nested_if, nullptr);
}

TEST_F(ParserBasicTest, ListLiteral) {
    viml::Parser parser("let list = [1, 2, \"hello\"]");
    auto program = parser.parse();
    
    auto let_stmt = dynamic_cast<viml::LetStatement*>(program->statements[0].get());
    auto list_literal = dynamic_cast<viml::ListLiteral*>(let_stmt->value.get());
    EXPECT_NE(list_literal, nullptr);
    EXPECT_EQ(list_literal->elements.size(), 3);
}

TEST_F(ParserBasicTest, DictLiteral) {
    viml::Parser parser("let dict = {\"name\": \"John\", \"age\": 30}");
    auto program = parser.parse();
    
    auto let_stmt = dynamic_cast<viml::LetStatement*>(program->statements[0].get());
    auto dict_literal = dynamic_cast<viml::DictLiteral*>(let_stmt->value.get());
    EXPECT_NE(dict_literal, nullptr);
    EXPECT_EQ(dict_literal->pairs.size(), 2);
}

TEST_F(ParserBasicTest, WhileStatement) {
    viml::Parser parser("while i < 10\nlet i = i + 1\nendwhile");
    auto program = parser.parse();
    
    auto while_stmt = dynamic_cast<viml::WhileStatement*>(program->statements[0].get());
    EXPECT_NE(while_stmt, nullptr);
    EXPECT_EQ(while_stmt->body.size(), 1);
}

TEST_F(ParserBasicTest, ForStatement) {
    viml::Parser parser("for item in [1, 2, 3]\necho item\nendfor");
    auto program = parser.parse();
    
    auto for_stmt = dynamic_cast<viml::ForStatement*>(program->statements[0].get());
    EXPECT_NE(for_stmt, nullptr);
    EXPECT_EQ(for_stmt->variable, "item");
    EXPECT_EQ(for_stmt->body.size(), 1);
}

} // namespace viml